#include <stdlib.h>
#include <stdio.h>
#include <time.h>

#include "lenet.h"
#include "load_model.h"

typedef struct MNIST {
    const char* train_image;
    const char* train_label;
    const char* test_image;
    const char* test_label;
} MNIST;

const int count_train=60000;
const int count_test=10000;
const char* model_file = "model.dat";

// dataset
int read_data(Image* data, uchar* label, const int count, const char* data_file, const char* label_file)
{
    FILE *fp_image = fopen(data_file, "rb");
    FILE *fp_label = fopen(label_file, "rb");
    if (!fp_image||!fp_label) {
        return 1;
    }
    fseek(fp_image, 16, SEEK_SET);
    fseek(fp_label, 8, SEEK_SET);
    fread(data, sizeof(*data)*count, 1, fp_image);
    fread(label,count, 1, fp_label);
    fclose(fp_image);
    fclose(fp_label);
    return 0;
}

// network
void training(LeNet5* lenet, Image* train_data, uchar* train_label, int batch_size, int total_size)
{
    for (int i = 0, percent = 0; i <= total_size - batch_size; i += batch_size)
    {
        TrainBatch(lenet, train_data + i, train_label + i, batch_size);
        if (i * 100 / total_size > percent) {
            printf("batchsize:%d\ttrain:%2d%%\n", batch_size, percent = i * 100 / total_size);
        }
    }
}

// network
int inference(LeNet5* lenet, Image* test_data, uchar* test_label, int total_size)
{
    int right = 0, percent = 0;
    for (int i = 0; i < total_size; ++i)
    {
        uchar l = test_label[i];
        int p = Predict(lenet, test_data[i], 10);
        right += l == p;
        if (i * 100 / total_size > percent) {
            printf("test:%2d%%\n", percent = i * 100 / total_size);
        }
    }
    return right;
}


// testbed
void run_lenet()
{
    Image *train_data = (Image *)calloc(count_train, sizeof(Image));
    uchar *train_label = (uchar *)calloc(count_train, sizeof(uchar));
    Image *test_data = (Image *)calloc(count_test, sizeof(Image));
    uchar *test_label = (uchar *)calloc(count_test, sizeof(uchar));

    MNIST mnist;
    mnist.train_image = "../mnist/train-images-idx3-ubyte";
    mnist.train_label = "../mnist/train-labels-idx1-ubyte";
    mnist.test_image = "../mnist/t10k-images-idx3-ubyte";
    mnist.test_label = "../mnist/t10k-labels-idx1-ubyte";

    if (read_data(train_data, train_label, count_train, mnist.train_image, mnist.train_label))
    {
        fprintf(stderr, "ERROR!!!\nDataset File Not Found!\n");
        free(train_data);
        free(train_label);
        exit(-1);
    }
    if (read_data(test_data, test_label, count_test, mnist.test_image, mnist.test_label))
    {
        fprintf(stderr, "ERROR!!!\nDataset File Not Found!\n");
        free(test_data);
        free(test_label);
        exit(-1);
    }

    LeNet5 *lenet = (LeNet5 *)malloc(sizeof(LeNet5));
    if (load_model(lenet, model_file)) {
        init_param(lenet);
        int num = 1;
        int* batches = (int*)malloc(sizeof(int)*num);
        for(int i=0; i<num; i++) {
            batches[i] = 300;
        }
        for (int i = 0; i < num; ++i) {
            training(lenet, train_data, train_label, batches[i], count_train);
        }
    }
    clock_t start = clock();
    int right = inference(lenet, test_data, test_label, count_test);
    printf("accuracy: %d/%d=%.2f%%\n", right, count_test, right*100.0/count_test);
    int time_cost = ( clock() - start ) / CLOCKS_PER_SEC;
    printf("Time: %ds\n", time_cost);
    save_model(lenet, model_file);
    free(lenet);
    free(train_data);
    free(train_label);
    free(test_data);
    free(test_label);
    printf("------\nDone\n");
}

int main()
{
    run_lenet();
    return 0;
}
